import os
import uuid
from itertools import groupby
from typing import List, Optional

import numpy as np

from common.logger import logger
from common.config import DEFAULT_CACHE_DIR
from common.typing_utils import numeric
from common.io_utils.dirs import mkdir_p
from common.io_utils.temps import temp_dir

from common.geometry.geom import Geom, BoundingBox
from common.geometry.tools import project_geom
from common.geoimage.gdal_sys import mosaic_ds
from common.geoimage.raster_dataset import RasterDataset

from catalog.models.base_stac_item import BaseSTACItem
from catalog.raster.utils import (
    determine_dest_srs,
    determine_dest_res,
    determine_dest_extent,
    is_s2_l2a_collection,
)

# NOTE: sentinel-2 L2A SCL data dictionary
# Sen2-Cor-L2A-Input-Output-Data-Definition-Document, Page 24-25, Table 8
# 0: nodata; 3: cloud shadow, 8: cloud medium prob, 9: cloud high prob, 11: snow or ice
# 10: thin cirrus, 7: unclassified, 6: water, 5: not veg, 4: veg, 2: dark features, shadows
TMP_DIR = os.path.join(DEFAULT_CACHE_DIR, "catalog", "raster")
mkdir_p(TMP_DIR)
DEFAULT_VAL_SCL_MASK = (0, 3, 8, 9, 11)


def composite_scl_mask_per_mgrs_tile(
    items, asset_names, aoi, vals_of_scl=DEFAULT_VAL_SCL_MASK
):

    if vals_of_scl is None:
        raise ValueError("Give at least 1 scl mask value")
    if len(asset_names) != 1 or asset_names[0] != "SCL":
        raise ValueError('Asset_names must be ["SCL"].')

    loaded_rds = []
    for item in items:
        vrt_scene = item.choose_assets(asset_names)
        rds = RasterDataset.from_vrt_scene(vrt_scene, aoi=aoi)
        loaded_rds.append(rds)

    result_mask = np.zeros((loaded_rds[0].data.shape))
    for rds_scl_tmp in loaded_rds:
        scl_mask = np.isin(np.squeeze(rds_scl_tmp.data), vals_of_scl)
        result_mask = np.logical_or(result_mask, ~scl_mask)

    result_mask = result_mask.astype(np.uint8)

    return RasterDataset.from_ndarray(result_mask, loaded_rds[-1].meta)


def composite_scl_mask_v1(
    items: List[BaseSTACItem],
    asset_names: List[str] = ["SCL"],
    extent: Optional[BoundingBox] = None,
    aoi: Optional[Geom] = None,
    dest_srs: Optional[str] = None,
    dest_res: Optional[numeric] = None,
    resample_alg: str = "near",
    **kwargs,
):
    """Composite a given list of `BaseSTACItems` calculating the valid mask after
    doing the medium composite.

    NOTE: This function is designed to be used within `catalog.raster.raster` only.
    It relies on the caller to provide several arguments correctly.
    """
    if not is_s2_l2a_collection(items):
        raise ValueError(
            (
                "The input items are not from the COPERNICUS/SENTINEL2/L2A product"
                " and not suppported by this composite strategy."
            )
        )

    if resample_alg != "near":
        raise ValueError("Resample algorithm must be near")

    if len(asset_names) != 1 or asset_names[0] != "SCL":
        raise ValueError('Asset_names must be ["SCL"]')

    vals_of_scl = kwargs["scl_val"] if "scl_val" in kwargs else DEFAULT_VAL_SCL_MASK

    if dest_srs is None:
        dest_srs = determine_dest_srs(aoi, items)

    if dest_res is None:
        dest_res = determine_dest_res(dest_srs, items)

    dest_extent = determine_dest_extent(extent, aoi, dest_srs, dest_res)

    # NOTE derive a `work_aoi`, a special part for the `s2decloud-v1` composite strategy
    # use that `work_aoi` to make per-tile process more efficient by only working on needed subset
    if dest_extent is not None:
        # if an extent is given, convert it to a geometry under WGS84 as the AOI
        bbox_geom = Geom.from_bbox(dest_extent)
        if dest_srs.upper() != "EPSG:4326":
            work_aoi = project_geom(bbox_geom, src_epsg=dest_srs, dest_epsg="EPSG:4326")
        else:
            work_aoi = bbox_geom
    else:
        work_aoi = None

    # group by mgrs tile
    items = sorted(items, key=lambda item: item.mgrs_tile)
    tile_groups = groupby(items, key=lambda item: item.mgrs_tile)
    n_tiles = len(set(item.mgrs_tile for item in items))

    if n_tiles == 1:
        for mgrs_tile, sub_items in tile_groups:
            sub_items = list(sub_items)
        logger.info(f"working on {mgrs_tile} with {len(sub_items)} images")
        rds = composite_scl_mask_per_mgrs_tile(
            sub_items, asset_names, work_aoi, vals_of_scl=vals_of_scl
        )
        rds = rds.warp(dest_srs=dest_srs, dest_res=dest_res, dest_extent=extent)
        return rds

    per_tile_files = []
    temp_dir_name = str(uuid.uuid4())
    with temp_dir(
        dir_name=temp_dir_name, base_dir=TMP_DIR, remove=True
    ) as temp_dir_path:
        for mgrs_tile, sub_items in tile_groups:
            sub_items = list(sub_items)
            logger.info(f"working on {mgrs_tile} with {len(sub_items)} images")
            rds = composite_scl_mask_per_mgrs_tile(
                sub_items, asset_names, work_aoi, vals_of_scl=vals_of_scl
            )
            if dest_srs != rds.meta.prj_epsg:
                rds = rds.warp(dest_srs=dest_srs)
            fpath_tmp = os.path.join(temp_dir_path, f"{mgrs_tile}.tif")
            rds.to_geotiff(fpath_tmp)
            per_tile_files.append(fpath_tmp)

        with mosaic_ds(
            per_tile_files,
            resample_alg=resample_alg,
            dest_res=dest_res,
            dest_extent=dest_extent,
            src_nodata=0,
        ) as gdal_rst:
            rds = RasterDataset.from_gdal_raster(gdal_rst, extent=dest_extent)

    return rds
